"""
身份证双面处理模块

专门处理身份证正面（人像面）和反面（国徽面）的组合排版
按照标准格式：人像面在上，国徽面在下
"""

import cv2
import numpy as np
import os
from typing import Tuple, Optional

from .image_matting import ImageMatting
from .image_correction import ImageCorrection
from .layout_generator import LayoutGenerator


class IDCardProcessor:
    """
    身份证双面处理器
    
    功能：
    - 处理身份证正面（人像面）
    - 处理身份证反面（国徽面）
    - 按标准格式组合：人像面在上，国徽面在下
    - 生成A4纸复印件效果
    """
    
    def __init__(self):
        """初始化身份证双面处理器"""
        self.matting = ImageMatting()
        self.correction = ImageCorrection()
        self.layout = LayoutGenerator()
        
        # 身份证双面排版参数
        self.front_back_spacing = 30  # 正反面间距（像素）
        self.page_margin = 50         # 页面边距
        
    def process_front_back(self, front_path: str, back_path: str, 
                          output_path: str, config: Optional[dict] = None) -> bool:
        """
        处理身份证正反面并组合排版
        
        参数:
            front_path: 正面（人像面）图片路径
            back_path: 反面（国徽面）图片路径
            output_path: 输出文件路径
            config: 处理配置
            
        返回:
            bool: 处理是否成功
            
        使用场景:
            processor = IDCardProcessor()
            success = processor.process_front_back("front.jpg", "back.jpg", "id_copy.jpg")
        """
        try:
            # 默认配置
            default_config = {
                'auto_correct': True,
                'enhance_contrast': True,
                'add_shadow': True,
                'scan_artifacts': True,
                'output_quality': 95
            }
            
            if config:
                default_config.update(config)
            
            print("开始处理身份证正反面...")
            
            # 步骤1: 检查输入文件
            if not os.path.exists(front_path):
                print(f"错误: 正面图片不存在 - {front_path}")
                return False
                
            if not os.path.exists(back_path):
                print(f"错误: 反面图片不存在 - {back_path}")
                return False
            
            # 步骤2: 处理正面（人像面）
            print("正在处理正面（人像面）...")
            front_processed = self._process_single_side(front_path, default_config)
            if front_processed is None:
                print("正面处理失败")
                return False
            
            # 步骤3: 处理反面（国徽面）
            print("正在处理反面（国徽面）...")
            back_processed = self._process_single_side(back_path, default_config)
            if back_processed is None:
                print("反面处理失败")
                return False
            
            # 步骤4: 组合正反面
            print("正在组合正反面...")
            combined_image = self._combine_front_back(front_processed, back_processed)
            
            # 步骤5: A4排版
            print("正在生成A4复印件...")
            a4_result = self._create_a4_layout(combined_image, default_config)
            
            # 步骤6: 保存结果
            print(f"正在保存到: {output_path}")
            self._save_image(a4_result, output_path, default_config['output_quality'])
            
            print("身份证双面处理完成!")
            return True
            
        except Exception as e:
            print(f"处理过程中发生错误: {e}")
            return False
    
    def _process_single_side(self, image_path: str, config: dict) -> Optional[np.ndarray]:
        """
        处理单面身份证
        
        参数:
            image_path: 图片路径
            config: 处理配置
            
        返回:
            numpy.ndarray: 处理后的图像，如果失败返回None
        """
        try:
            # 抠图去背景
            rgba_image = self.matting.remove_background(image_path)
            
            # 添加白色背景
            rgb_image = self.matting.add_white_background(rgba_image)
            
            # 图像校正
            if config['auto_correct']:
                rgb_image = self.correction.auto_correct(rgb_image)
            elif config['enhance_contrast']:
                rgb_image = self.correction.enhance_contrast(rgb_image)
            
            # 调整为身份证标准尺寸
            id_sized = self.layout.resize_to_id_card(rgb_image)
            
            return id_sized
            
        except Exception as e:
            print(f"单面处理失败: {e}")
            return None
    
    def _combine_front_back(self, front_image: np.ndarray, 
                           back_image: np.ndarray) -> np.ndarray:
        """
        组合正反面图像（人像面在上，国徽面在下）
        
        参数:
            front_image: 正面图像
            back_image: 反面图像
            
        返回:
            numpy.ndarray: 组合后的图像
        """
        try:
            # 确保两张图片尺寸一致
            front_h, front_w = front_image.shape[:2]
            back_h, back_w = back_image.shape[:2]
            
            # 使用较大的宽度，确保对齐
            target_width = max(front_w, back_w)
            
            # 调整图片尺寸
            if front_w != target_width:
                scale = target_width / front_w
                new_height = int(front_h * scale)
                front_image = cv2.resize(front_image, (target_width, new_height))
                front_h = new_height
            
            if back_w != target_width:
                scale = target_width / back_w
                new_height = int(back_h * scale)
                back_image = cv2.resize(back_image, (target_width, new_height))
                back_h = new_height
            
            # 计算组合图像尺寸
            combined_height = front_h + back_h + self.front_back_spacing
            combined_width = target_width
            
            # 创建组合画布（白色背景）
            combined = np.ones((combined_height, combined_width, 3), dtype=np.uint8) * 255
            
            # 放置正面（人像面）在上
            combined[0:front_h, 0:target_width] = front_image
            
            # 放置反面（国徽面）在下
            y_offset = front_h + self.front_back_spacing
            combined[y_offset:y_offset+back_h, 0:target_width] = back_image
            
            return combined
            
        except Exception as e:
            print(f"组合图像失败: {e}")
            # 如果组合失败，返回正面图像
            return front_image
    
    def _create_a4_layout(self, combined_image: np.ndarray, config: dict) -> np.ndarray:
        """
        创建A4页面布局
        
        参数:
            combined_image: 组合后的身份证图像
            config: 配置参数
            
        返回:
            numpy.ndarray: A4布局的图像
        """
        try:
            # 创建A4背景
            a4_height, a4_width = self.layout.a4_height, self.layout.a4_width
            a4_background = np.ones((a4_height, a4_width, 3), dtype=np.uint8) * 248
            
            # 计算居中位置
            img_h, img_w = combined_image.shape[:2]
            
            # 确保图像不超出A4页面
            max_width = a4_width - 2 * self.page_margin
            max_height = a4_height - 2 * self.page_margin
            
            if img_w > max_width or img_h > max_height:
                # 计算缩放比例
                scale_w = max_width / img_w
                scale_h = max_height / img_h
                scale = min(scale_w, scale_h)
                
                # 缩放图像
                new_width = int(img_w * scale)
                new_height = int(img_h * scale)
                combined_image = cv2.resize(combined_image, (new_width, new_height))
                img_h, img_w = new_height, new_width
            
            # 计算居中位置
            x_pos = (a4_width - img_w) // 2
            y_pos = (a4_height - img_h) // 2
            
            # 直接放置图像（双面组合已经包含了阴影效果逻辑）
            a4_background[y_pos:y_pos+img_h, x_pos:x_pos+img_w] = combined_image
            
            # 添加扫描效果
            if config.get('scan_artifacts', True):
                a4_background = self.layout.add_scan_artifacts(a4_background)
            
            return a4_background
            
        except Exception as e:
            print(f"A4布局创建失败: {e}")
            return combined_image
    
    def _save_image(self, image: np.ndarray, output_path: str, quality: int = 95):
        """
        保存图像
        
        参数:
            image: 图像数组
            output_path: 输出路径
            quality: 图像质量
        """
        try:
            # 确保输出目录存在
            output_dir = os.path.dirname(output_path)
            if output_dir:
                os.makedirs(output_dir, exist_ok=True)
            
            # 转换颜色空间
            if len(image.shape) == 3 and image.shape[2] == 3:
                image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
            else:
                image_bgr = image
            
            # 保存图像
            cv2.imwrite(output_path, image_bgr, 
                       [cv2.IMWRITE_JPEG_QUALITY, quality])
            
        except Exception as e:
            print(f"保存图像失败: {e}")
    
    def get_processing_info(self, front_path: str, back_path: str) -> dict:
        """
        获取双面处理信息
        
        参数:
            front_path: 正面图片路径
            back_path: 反面图片路径
            
        返回:
            dict: 处理信息
        """
        try:
            info = {
                "front_image": "未找到",
                "back_image": "未找到",
                "processing_steps": [
                    "正面抠图和校正",
                    "反面抠图和校正", 
                    "双面组合排版",
                    "A4页面布局",
                    "扫描效果处理"
                ]
            }
            
            # 检查正面
            if os.path.exists(front_path):
                front_img = cv2.imread(front_path)
                if front_img is not None:
                    h, w = front_img.shape[:2]
                    size = os.path.getsize(front_path)
                    info["front_image"] = f"{w}x{h} ({size/1024:.1f}KB)"
            
            # 检查反面
            if os.path.exists(back_path):
                back_img = cv2.imread(back_path)
                if back_img is not None:
                    h, w = back_img.shape[:2]
                    size = os.path.getsize(back_path)
                    info["back_image"] = f"{w}x{h} ({size/1024:.1f}KB)"
            
            return info
            
        except Exception as e:
            return {"error": f"获取信息失败: {e}"}
